import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error
from skopt import gp_minimize
from skopt.space import Real
from typing import Tuple, List, Dict

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

# ---------------------------
# 增强版数据预处理
# ---------------------------
def preprocess_data(df: pd.DataFrame) -> pd.DataFrame:
    """处理原始数据，提取男胎样本并计算达标时间
    
    Args:
        df: 包含NIPT检测数据的DataFrame
        
    Returns:
        处理后的DataFrame，仅包含男胎样本和达标时间
    """
    # 检查必要列是否存在
    required_columns = ['孕妇代码', 'Y染色体浓度', '检测孕周']
    if not all(col in df.columns for col in required_columns):
        raise ValueError("输入数据缺少必要的列：'孕妇代码', 'Y染色体浓度', '检测孕周'")
    
    # 增强版孕周转换函数
    def convert_gestational_week(week_str):
        try:
            if isinstance(week_str, (int, float)):
                return float(week_str)
            if '+' in week_str:
                w, d = week_str.lower().replace('w', '').split('+')
                return float(w) + float(d)/7
            return float(week_str.lower().replace('w', ''))
        except:
            return np.nan
    
    df['检测孕周'] = df['检测孕周'].apply(convert_gestational_week)
    df = df.dropna(subset=['检测孕周'])
    
    # 筛选男胎数据（Y染色体浓度≥4%为男胎）
    male_data = df[df['Y染色体浓度'] >= 0.04].copy()
    
    # 更高效地计算达标孕周
    male_data['达标孕周'] = male_data.groupby('孕妇代码')['检测孕周'].transform('min')
    male_data = male_data.dropna(subset=['达标孕周'])
    
    
    return male_data

# ---------------------------
# 增强版预测模型构建
# ---------------------------
def build_time_prediction_model(data: pd.DataFrame) -> Tuple[RandomForestRegressor, PolynomialFeatures, StandardScaler]:
    """建立达标时间预测模型
    
    Args:
        data: 预处理后的数据
        
    Returns:
        model: 训练好的随机森林模型
        poly: 多项式特征转换器
        scaler: 标准化转换器
    """
    # 检查必要特征列是否存在
    features = ['年龄', '身高', '体重', '孕妇BMI', '怀孕次数', '生产次数']
    missing_features = [f for f in features if f not in data.columns]
    if missing_features:
        raise ValueError(f"输入数据缺少必要的特征列：{missing_features}")
    
    X = data[features]
    y = data['达标孕周']
    
    # 添加多项式特征和标准化
    poly = PolynomialFeatures(degree=2, include_bias=False)
    scaler = StandardScaler()
    X_poly = scaler.fit_transform(poly.fit_transform(X))
    
    # 训练随机森林模型
    X_train, X_test, y_train, y_test = train_test_split(
        X_poly, y, test_size=0.2, random_state=42
    )
    
    model = RandomForestRegressor(
        n_estimators=150,
        max_depth=12,
        min_samples_split=5,
        random_state=42,
        n_jobs=-1
    )
    model.fit(X_train, y_train)
    
    # 更全面的模型评估
    train_pred = model.predict(X_train)
    test_pred = model.predict(X_test)
    
    print("\n=== 模型评估 ===")
    print(f"训练集R²: {model.score(X_train, y_train):.3f}")
    print(f"测试集R²: {model.score(X_test, y_test):.3f}")
    print(f"训练集MAE: {mean_absolute_error(y_train, train_pred):.3f} 周")
    print(f"测试集MAE: {mean_absolute_error(y_test, test_pred):.3f} 周")
    print(f"训练集RMSE: {np.sqrt(mean_squared_error(y_train, train_pred)):.3f} 周")
    print(f"测试集RMSE: {np.sqrt(mean_squared_error(y_test, test_pred)):.3f} 周")
    
    return model, poly, scaler

# ---------------------------
# 优化算法实现（修复了比较错误）
# ---------------------------
def optimize_bmi_groups(data: pd.DataFrame, max_groups: int = 4, min_samples: int = 30) -> List[float]:
    """动态规划优化BMI分组
    
    Args:
        data: 预处理后的数据
        max_groups: 最大分组数
        min_samples: 每组最小样本数
        
    Returns:
        分组边界值列表
    """
    data = data.sort_values('孕妇BMI').reset_index(drop=True)
    n = len(data)
    
    # 数据量校验
    if n < min_samples * max_groups:
        raise ValueError(f"数据量不足，至少需要{min_samples*max_groups}个样本")
    
    # DP表初始化
    dp = [[float('inf')] * (n+1) for _ in range(max_groups+1)]
    dp[0][0] = 0
    splits = [[[] for _ in range(n+1)] for _ in range(max_groups+1)]
    
    # 预处理风险矩阵（限制分组大小）
    risk_matrix = np.zeros((n, n))
    for i in range(n):
        for j in range(i, min(i + 50, n)):  # 限制最大分组大小50
            if j - i + 1 >= min_samples:  # 满足最小样本量
                group_data = data.iloc[i:j+1]
                risk_matrix[i][j] = calculate_group_risk(group_data)
    
    # DP填充（限制回溯范围）
    for k in range(1, max_groups+1):
        for i in range(1, n+1):
            min_j = max(k-1, i - 50)  # 限制回溯范围
            for j in range(min_j, i):
                if dp[k-1][j] + risk_matrix[j][i-1] < dp[k][i]:
                    dp[k][i] = dp[k-1][j] + risk_matrix[j][i-1]
                    splits[k][i] = splits[k-1][j] + [j]
    
    # 提取最优分割点
    best_splits = splits[max_groups][n]
    bounds = sorted(list(set([data['孕妇BMI'].min()] + 
                            data.iloc[best_splits]['孕妇BMI'].tolist() + 
                            [data['孕妇BMI'].max()])))
    
    return bounds

def calculate_group_risk(group_data: pd.DataFrame) -> float:
    """计算单组的风险值
    
    Args:
        group_data: 单个BMI分组的数据
        
    Returns:
        该组的综合风险值
    """
    def time_objective(t: float) -> float:
        # 修复比较错误：确保t是标量值
        t_value = t[0] if isinstance(t, (list, np.ndarray)) else t
        pred_prob = predict_detection_probability(group_data, t_value)
        risk_weight = 1 if t_value <= 12 else (2 if t_value <= 27 else 5)
        return -pred_prob * risk_weight  # 最小化
    
    # 增加优化迭代次数
    res = gp_minimize(
        time_objective,
        [Real(8, 40)],
        n_calls=50,
        random_state=42,
        acq_func='EI',
        verbose=False
    )
    return -res.fun

def predict_detection_probability(group_data: pd.DataFrame, t: float) -> float:
    """预测在时间t的达标概率"""
    # 修复比较错误：确保t是标量值
    t_value = t[0] if isinstance(t, (list, np.ndarray)) else t
    early_records = group_data[group_data['检测孕周'] <= t_value]
    return (early_records['Y染色体浓度'] >= 0.04).mean() if len(early_records) > 0 else 0.0

# ---------------------------
# 增强版主流程
# ---------------------------
def main(df: pd.DataFrame) -> pd.DataFrame:
    """主优化流程
    
    Args:
        df: 原始NIPT检测数据
        
    Returns:
        分组优化结果DataFrame
    """
    print("\n=== 数据预处理 ===")
    male_data = preprocess_data(df)
    print(f"有效男胎样本数: {len(male_data)}")
    
    print("\n=== 模型构建 ===")
    time_model, poly, scaler = build_time_prediction_model(male_data)
    
    print("\n=== BMI分组优化 ===")
    bmi_bounds = optimize_bmi_groups(male_data)
    print(f"获得BMI分组边界: {bmi_bounds}")
    
    # 为每个分组确定最佳检测时间
    groups = []
    for i in range(len(bmi_bounds)-1):
        lower, upper = bmi_bounds[i], bmi_bounds[i+1]
        group_data = male_data[
            (male_data['孕妇BMI'] >= lower) & 
            (male_data['孕妇BMI'] < upper)
        ]
        
        if len(group_data) < 30:  # 忽略小样本组
            continue
            
        res = gp_minimize(
            lambda t: calculate_group_risk(group_data),
            [Real(8, 40)],
            n_calls=50,
            random_state=42,
            acq_func='EI',
            verbose=False
        )
        best_time = round(res.x[0], 1)
        
        groups.append({
            'BMI范围': f"{lower:.1f}-{upper:.1f}",
            '最佳检测孕周': best_time,
            '样本量': len(group_data),
            '预期达标率': round(predict_detection_probability(group_data, best_time), 3),
            '风险等级': '低' if best_time <=12 else ('中' if best_time <=27 else '高')
        })
    
    # 结果展示
    results = pd.DataFrame(groups)
    print("\n=== 最佳分组策略 ===")
    print(results.to_markdown(index=False))
    
    # 增强可视化
    plot_results(results)
    
    return results

def plot_results(results: pd.DataFrame) -> None:
    """增强版可视化结果"""
    plt.figure(figsize=(14, 7))
    
    # 创建颜色映射和样式
    color_map = {'低': '#4CAF50', '中': '#FFC107', '高': '#F44336'}
    bar_width = 0.6
    
    # 绘制柱状图
    for _, group in results.iterrows():
        plt.bar(
            group['BMI范围'], 
            group['最佳检测孕周'],
            width=bar_width,
            label=f"{group['风险等级']}风险",
            color=color_map[group['风险等级']],
            alpha=0.8,
            edgecolor='black'
        )
        
        # 添加数据标签
        plt.text(
            group['BMI范围'], 
            group['最佳检测孕周'] + 0.3,
            f"{group['最佳检测孕周']}周\n达标率:{group['预期达标率']*100:.1f}%",
            ha='center',
            va='bottom',
            fontsize=10,
            bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5)
        )
    
    # 美化图表
    plt.xlabel('BMI分组', fontsize=12, labelpad=10)
    plt.ylabel('推荐检测孕周', fontsize=12, labelpad=10)
    plt.title('基于BMI的最佳NIPT检测时间推荐', fontsize=14, pad=20, fontweight='bold')
    
    # 调整图例和网格
    plt.legend(title='风险等级', loc='upper right', framealpha=0.9)
    plt.grid(axis='y', linestyle='--', alpha=0.5)
    
    # 调整y轴范围
    y_max = results['最佳检测孕周'].max() + 2
    plt.ylim(8, y_max if y_max > 20 else 20)
    
    plt.tight_layout()
    plt.savefig('ques3_optimal_bmi_groups.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("\n已保存可视化结果: ques3_optimal_bmi_groups.png")

# ---------------------------
# 敏感性分析
# ---------------------------
def sensitivity_analysis(results: pd.DataFrame, male_data: pd.DataFrame, 
                       error_rates: List[float] = [0.01, 0.05, 0.1]) -> pd.DataFrame:
    """分析检测误差对结果的影响
    
    Args:
        results: 主流程的结果DataFrame
        male_data: 预处理后的男胎数据
        error_rates: 要分析的误差率列表
        
    Returns:
        敏感性分析结果DataFrame
    """
    sensitivity_results = []
    
    for error_rate in error_rates:
        for _, group in results.iterrows():
            lower, upper = map(float, group['BMI范围'].split('-'))
            group_data = male_data[
                (male_data['孕妇BMI'] >= lower) & 
                (male_data['孕妇BMI'] < upper)
            ]
            
            if len(group_data) == 0:
                continue
                
            # 模拟误差影响：达标率按误差率下降
            original_rate = predict_detection_probability(group_data, group['最佳检测孕周'])
            adjusted_rate = original_rate * (1 - error_rate)
            
            sensitivity_results.append({
                '误差率': f"{error_rate*100:.0f}%",
                'BMI范围': group['BMI范围'],
                '原始达标率': round(original_rate, 3),
                '调整后达标率': round(adjusted_rate, 3),
                '达标率损失': round(original_rate - adjusted_rate, 3)
            })
    
    sensitivity_df = pd.DataFrame(sensitivity_results)
    
    # 可视化
    plt.figure(figsize=(14, 7))
    for error_rate in error_rates:
        subset = sensitivity_df[sensitivity_df['误差率'] == f"{error_rate*100:.0f}%"]
        plt.plot(subset['BMI范围'], subset['调整后达标率'], 
                marker='o', markersize=8, linewidth=2,
                label=f'误差率={error_rate*100:.0f}%')
    
    plt.xlabel('BMI分组', fontsize=12)
    plt.ylabel('调整后达标率', fontsize=12)
    plt.title('检测误差对达标率的影响', fontsize=14, pad=20)
    plt.legend(title='误差率', loc='upper right')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.ylim(0, 1)
    plt.tight_layout()
    plt.savefig('sensitivity_analysis.png', dpi=300)
    plt.close()
    print("\n已保存敏感性分析结果: sensitivity_analysis.png")
    
    return sensitivity_df

# ---------------------------
# 执行主程序
# ---------------------------
if __name__ == "__main__":
    try:
        print("=== 数据加载 ===")
        df = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')
        
        print("\n=== 主优化流程 ===")
        optimal_groups = main(df)
        
        print("\n=== 敏感性分析 ===")
        male_data = preprocess_data(df)
        sensitivity_results = sensitivity_analysis(optimal_groups, male_data)
        print("\n敏感性分析结果：")
        print(sensitivity_results.to_markdown(index=False))
        
    except Exception as e:
        print(f"程序运行出错: {str(e)}")
        raise